package com.ruoyi.huanjing.tool;

import ucar.nc2.NetcdfFile;
import ucar.nc2.Variable;
import ucar.nc2.dataset.NetcdfDataset;
import ucar.ma2.Array;
import ucar.ma2.Index;
import org.jfree.chart.ChartFactory;
import org.jfree.chart.ChartFrame;
import org.jfree.chart.JFreeChart;
import org.jfree.chart.plot.PlotOrientation;
import org.jfree.chart.plot.XYPlot;
import org.jfree.chart.renderer.xy.XYBlockRenderer;
import org.jfree.data.xy.DefaultXYZDataset;
import org.jfree.data.xy.XYDataset;
import org.jfree.data.xy.XYZDataset;
import javax.swing.*;
import java.awt.*;
import java.io.IOException;
import java.util.Arrays;
import java.util.List;

public class HDF4HeatmapPlotter {

    public static void main(String[] args) {
        String filePath = "G:\\MOD11C1\\MOD11C1.A2004001.061.2020168035507.hdf";
        String varName = "MODIS_CMG_3MIN_LST/Data_Fields/QC_Day";

        try {
            plotGeoHeatmap(filePath, varName, "Lat", "Lon");
        } catch (IOException e) {
            e.printStackTrace();
            JOptionPane.showMessageDialog(null, "Error: " + e.getMessage(),
                    "Plotting Error", JOptionPane.ERROR_MESSAGE);
        }
    }

    public static void plotGeoHeatmap(String filePath, String varName,
                                      String latVarName, String lonVarName) throws IOException {

        System.out.println("Opening HDF4 file: " + filePath);
        try (NetcdfFile ncFile = NetcdfDataset.openFile(filePath, null)) {

            System.out.println("Successfully opened file. All variables: " +
                    Arrays.toString(ncFile.getVariables().toArray()));

            // Check if variable exists
            Variable var = ncFile.findVariable(varName);
            if (var == null) {
                throw new IllegalArgumentException("Variable " + varName +
                        " not found in file. Available variables: " +
                        Arrays.toString(ncFile.getVariables().toArray()));
            }

            System.out.println("Variable " + varName + " shape: " +
                    Arrays.toString(var.getShape()));

            // Read variable data
            Array varData = var.read();

            // Get latitude and longitude variables
            Variable latVar = ncFile.findVariable(latVarName);
            Variable lonVar = ncFile.findVariable(lonVarName);

            if (latVar == null || lonVar == null) {
                // Try to find them automatically
                List<String> latVarNames = Arrays.asList("lat", "Latitude", "Lat", "y", "Y");
                List<String> lonVarNames = Arrays.asList("lon", "Longitude", "Lon", "x", "X");

                System.out.println("Trying to find lat/lon variables automatically...");

                for (String name : latVarNames) {
                    latVar = ncFile.findVariable(name);
                    if (latVar != null) {
                        System.out.println("Found latitude variable: " + name);
                        break;
                    }
                }

                for (String name : lonVarNames) {
                    lonVar = ncFile.findVariable(name);
                    if (lonVar != null) {
                        System.out.println("Found longitude variable: " + name);
                        break;
                    }
                }
            }

            // Determine variable type
            String varType;
            boolean isGeo2D = false;

            if (latVar != null && lonVar != null) {
                isGeo2D = true;
                varType = "Geo2D";
                System.out.println("Identified as geographic data");
            } else if (var.getRank() == 1) {
                varType = "1D";
            } else if (var.getRank() == 2) {
                varType = "2D";
                System.out.println("Identified as regular 2D data");
            } else {
                // Handle higher dimensions - take first 2D slice
                if (var.getRank() > 2) {
                    System.out.println("Higher dimensional data - taking first 2D slice");
                    varType = "2D";
                } else {
                    throw new IllegalArgumentException("Variable " + varName +
                            " has unsupported dimensions");
                }
            }

            System.out.println("Final variable type: " + varType);

            // Create and show appropriate plot
            if (varType.equals("1D")) {
                plot1DData(varData, varName);
            } else if (varType.equals("2D") || varType.equals("Geo2D")) {
                if (isGeo2D) {
                    Array latData = latVar.read();
                    Array lonData = lonVar.read();
                    plotGeoHeatmap(lonData, latData, varData, varName);
                } else {
                    plot2DHeatmap(varData, varName);
                }
            }
        }
    }

    private static void plot1DData(Array data, String title) {
        double[] x = new double[data.getShape()[0]];
        double[] y = new double[data.getShape()[0]];

        Index index = data.getIndex();
        for (int i = 0; i < data.getShape()[0]; i++) {
            x[i] = i;
            y[i] = data.getDouble(index.set(i));
        }

        JFreeChart chart = ChartFactory.createXYLineChart(
                title + " - Line Plot",
                "Index",
                "Value",
                createXYDataset(x, y),
                PlotOrientation.VERTICAL,
                true,
                true,
                false
        );

        displayChart(chart);
    }

    private static void plot2DHeatmap(Array data, String title) {
        int[] shape = data.getShape();
        double[] x = new double[shape[0] * shape[1]];
        double[] y = new double[shape[0] * shape[1]];
        double[] z = new double[shape[0] * shape[1]];

        Index index = data.getIndex();
        int idx = 0;
        for (int i = 0; i < shape[0]; i++) {
            for (int j = 0; j < shape[1]; j++) {
                x[idx] = j;
                y[idx] = i;
                z[idx] = data.getDouble(index.set(i, j));
                idx++;
            }
        }

        XYZDataset dataset = createXYZDataset(x, y, z);
        JFreeChart chart = createHeatmapChart(dataset, title + " - Heatmap", "X", "Y");
        displayChart(chart);
    }

    private static void plotGeoHeatmap(Array lonData, Array latData, Array varData, String title) {
        // For simplicity, we'll just plot as regular heatmap here
        // A proper geographic plot would require more complex libraries like GeoTools

        int[] shape = varData.getShape();
        double[] x = new double[shape[0] * shape[1]];
        double[] y = new double[shape[0] * shape[1]];
        double[] z = new double[shape[0] * shape[1]];

        Index varIndex = varData.getIndex();
        Index latIndex = latData.getIndex();
        Index lonIndex = lonData.getIndex();

        int idx = 0;
        for (int i = 0; i < shape[0]; i++) {
            for (int j = 0; j < shape[1]; j++) {
                x[idx] = lonData.getDouble(lonIndex.set(j)); // Assuming lon is 1D
                y[idx] = latData.getDouble(latIndex.set(i)); // Assuming lat is 1D
                z[idx] = varData.getDouble(varIndex.set(i, j));
                idx++;
            }
        }

        XYZDataset dataset = createXYZDataset(x, y, z);
        JFreeChart chart = createHeatmapChart(dataset, title + " - Geographic Heatmap",
                "Longitude", "Latitude");
        displayChart(chart);
    }

    private static XYZDataset createXYZDataset(double[] x, double[] y, double[] z) {
        DefaultXYZDataset dataset = new DefaultXYZDataset();
        double[][] data = new double[3][x.length];
        System.arraycopy(x, 0, data[0], 0, x.length);
        System.arraycopy(y, 0, data[1], 0, y.length);
        System.arraycopy(z, 0, data[2], 0, z.length);
        dataset.addSeries("Data", data);
        return dataset;
    }

    private static XYDataset createXYDataset(double[] x, double[] y) {
        DefaultXYZDataset dataset = new DefaultXYZDataset();
        double[][] data = new double[3][x.length];
        System.arraycopy(x, 0, data[0], 0, x.length);
        System.arraycopy(y, 0, data[1], 0, y.length);
        Arrays.fill(data[2], 1.0); // Dummy Z values
        dataset.addSeries("Data", data);
        return dataset;
    }

    private static JFreeChart createHeatmapChart(XYZDataset dataset, String title,
                                                 String xLabel, String yLabel) {
        JFreeChart chart = ChartFactory.createScatterPlot(
                title,
                xLabel,
                yLabel,
                dataset,
                PlotOrientation.VERTICAL,
                true,
                true,
                false
        );

        XYPlot plot = (XYPlot) chart.getPlot();
        XYBlockRenderer renderer = new XYBlockRenderer();
        plot.setRenderer(renderer);

        // Configure the renderer
        renderer.setPaintScale(new org.jfree.chart.renderer.LookupPaintScale(
                getMinValue(dataset), getMaxValue(dataset),
                new Color(0, 0, 0, 0))); // Default to transparent

        // You can customize the color scale here

        return chart;
    }

    private static double getMinValue(XYZDataset dataset) {
        double min = Double.MAX_VALUE;
        for (int i = 0; i < dataset.getItemCount(0); i++) {
            double z = dataset.getZValue(0, i);
            if (z < min) min = z;
        }
        return min;
    }

    private static double getMaxValue(XYZDataset dataset) {
        double max = Double.MIN_VALUE;
        for (int i = 0; i < dataset.getItemCount(0); i++) {
            double z = dataset.getZValue(0, i);
            if (z > max) max = z;
        }
        return max;
    }

    private static void displayChart(JFreeChart chart) {
        ChartFrame frame = new ChartFrame("Data Visualization", chart);
        frame.setSize(800, 600);
        frame.setVisible(true);
    }
}